import sys
import torch
import torch.nn as nn
import matplotlib as mpl
import matplotlib.pyplot as plt
import networkx
import NeuronUnit_2025_3_12 as nu_group

def print_progress_bar(iteration, total, prefix='', length=50):
    percent = (iteration / total) * 100
    filled_length = int(length * iteration // total)
    bar = '█' * filled_length + '-' * (length - filled_length)
    sys.stdout.write(f'\r |{bar}|  {'iteration='}{prefix}  {percent:.2f}%')
    sys.stdout.flush()

class FreelyEmergedNeuralNetworks():
    #主要功能函数
    def __init__(self,connect_num_linear=4,connect_num_exponent=2,turning_point_num=5,
                 value_memory_num=5,net_scale=10,entry_num=2,dropout_rate=0.2,
                 device=torch.device("cuda" if torch.cuda.is_available() else "cpu")):     #网络初始化函数
        self.nu_group_registration=[[False,entry_num],[True,net_scale]]
        self.neuron_unit_group=[nu_group.NeuronUnitGroup(connect_num_linear=connect_num_linear,
                                                connect_num_exponent=connect_num_exponent,
                                                turning_point_num=turning_point_num,
                                                value_memory_num=value_memory_num,group_scale=net_scale,
                                                last_frame_size=net_scale+entry_num,device=device)]
        self.last_frame_size=net_scale+entry_num
        self.dropout=nn.Dropout(dropout_rate)
    def compute(self,input_tensor_list):
        last_frame=torch.tensor([])
        i_e=0
        i_c=0
        for registration in self.nu_group_registration:
            if registration[0]:
                last_frame=torch.cat((last_frame,self.neuron_unit_group[i_c].value))
                i_c+=1
            else:
                last_frame=torch.cat((last_frame,torch.flatten(input_tensor_list[i_e])))
                i_e+=1
        last_frame=self.dropout(last_frame)
        for i in range(len(self.neuron_unit_group)):
            self.neuron_unit_group[i].compute(last_frame)
    def before_forget(self,forget_rate):
        for i in range(len(self.neuron_unit_group)):
            self.neuron_unit_group[i].before_forget(forget_rate)
    def forward(self,rounds,input_tensor_list_all_rounds,forget_rate=1.):
        for i in range(len(self.neuron_unit_group)):
            self.neuron_unit_group[i].forget(forget_rate)
        for i in range(rounds):
            self.compute(input_tensor_list_all_rounds[i])
            print_progress_bar(i+1,rounds,i+1)
    def output_results(self,sampling_list): #=[[0,[0,1,2,3]] , [1,[0,1]]]
        results=torch.tensor([])
        for list in sampling_list:
            results=torch.cat((results,self.neuron_unit_group[list[0]].value[list[1]]))
        return results
    def restart(self):
        for i in range(len(self.neuron_unit_group)):
            self.neuron_unit_group[i].restart()
    def reload(self):
        for i in range(len(self.neuron_unit_group)):
            self.neuron_unit_group[i].reload()
    def pruning(self,level_linear=0.001,level_exponent=0.001):
        for i in range(len(self.neuron_unit_group)):
            self.neuron_unit_group[i].pruning(level_linear=level_linear,level_exponent=level_exponent)
    def supconnect(self,num_linear=1,num_exponent=1):
        for i in range(len(self.neuron_unit_group)):
            self.neuron_unit_group[i].supconnect(last_frame_size=self.last_frame_size,num_linear=num_linear,num_exponent=num_exponent)
    def scaling(self,is_nu_group,connect_num_linear=4,connect_num_exponent=2,turning_point_num=5,
                 value_memory_num=5,adding_group_scale=10,
                 device=torch.device("cuda" if torch.cuda.is_available() else "cpu")):
        self.nu_group_registration.append([is_nu_group,adding_group_scale])
        self.last_frame_size+=adding_group_scale
        if is_nu_group:
            self.neuron_unit_group+=[nu_group.NeuronUnitGroup(connect_num_linear=connect_num_linear,
                                                connect_num_exponent=connect_num_exponent,
                                                turning_point_num=turning_point_num,
                                                value_memory_num=value_memory_num,group_scale=adding_group_scale,
                                                last_frame_size=self.last_frame_size,device=device)]
    def unforget(self,forget_rate):
        for i in range(len(self.neuron_unit_group)):
            self.neuron_unit_group[i](forget_rate)
    def parameters(self):
        parameters=[]
        for nu_group in self.neuron_unit_group:
            parameters.append(nu_group.parameters())
        return parameters
    def load_parameters(self,file_name='parameters.pth'):
        parameters=torch.load(file_name)
        for i in range(len(self.neuron_unit_group)):
            self.neuron_unit_group[i].load_parameters(parameters[i])
    def save_parameters(self,file_name='parameters.pth'):
        parameters=self.parameters()
        torch.save(parameters,file_name)
    def param_4_optim(self):
        param_4_optim=[]
        for i in range(len(self.neuron_unit_group)):
            param_4_optim+=self.neuron_unit_group[i].param_4_optim()
        return param_4_optim
    def frame(self):
        frame=[]
        for nu_group in self.neuron_unit_group:
            frame.append(nu_group.frame())
        return frame
    def load_frame(self,frame):
        for i in range(len(self.neuron_unit_group)):
            self.neuron_unit_group[i].load_frame(frame[i])

    #实验研究用函数
    def activation_demostrate(self,nu_position_list):   #[1,2]
        x=[i/100-3 for i in range(600)]
        y=[self.neuron_unit_group[nu_position_list[0]][nu_position_list[1]](i).item() for i in x]
        ax=plt.gca()
        ax.spines['right'].set_color('none')
        ax.spines['top'].set_color('none')
        ax.xaxis.set_ticks_position('bottom')
        ax.spines['bottom'].set_position(('data',0))
        ax.yaxis.set_ticks_position('left')
        ax.spines['left'].set_position(('data',0))
        plt.plot(x,y)
        plt.show()
    def graph_paint(self):
        edges=[]
        nu_num0=0
        nu_num1=0
        nu_group=0
        for registration in self.nu_group_registration:
            if registration[0]:
                for nu in self.neuron_unit_group[nu_group].connect_linear:
                    for i in range(len(nu[0])):
                        edges.append([nu[0][i],nu_num0,nu[1][i].item()])
                    nu_num0+=1
                for nu in self.neuron_unit_group[nu_group].connect_exponent:
                    for i in range(len(nu[0])):
                        edges.append([nu[0][i],nu_num1,nu[1][i].item()])
                    nu_num1+=1
                    print_progress_bar(nu_num1,self.last_frame_size,nu_num1)
                nu_group+=1
            else:
                nu_num0+=registration[1]
                nu_num1+=registration[1]
                print_progress_bar(nu_num1,self.last_frame_size,nu_num1)
        self.G=networkx.DiGraph()
        self.G.clear()
        self.G.add_weighted_edges_from(edges)
        self.pos=networkx.spring_layout(self.G)
    def graph_demostrate(self,width=0.1,node_size=15,font_size=6,figsize=(19.2,10.8)):
        plt.figure(figsize=figsize)
        edges=[]
        nu_num0=0
        nu_num1=0
        nu_group=0
        entry_node_list=[]
        for registration in self.nu_group_registration:
            if registration[0]:
                for nu in self.neuron_unit_group[nu_group].connect_linear:
                    for i in range(len(nu[0])):
                        edges.append([nu[0][i],nu_num0,nu[1][i].item()])
                    nu_num0+=1
                for nu in self.neuron_unit_group[nu_group].connect_exponent:
                    for i in range(len(nu[0])):
                        edges.append([nu[0][i],nu_num1,nu[1][i].item()])
                    nu_num1+=1
                nu_group+=1
            else:
                entry_node_list+=[i for i in range(nu_num0,nu_num0+registration[1])]
                nu_num0+=registration[1]
                nu_num1+=registration[1]
        self.G.clear()
        self.G.add_weighted_edges_from(edges)
        w=[abs(i[2]) for i in edges]
        edges1=networkx.draw_networkx_edges(self.G,self.pos,node_size=node_size,arrowstyle="->",arrowsize=width*10,edge_color='#245DA2',edge_cmap=plt.cm.Blues,width=width,)
        pc=mpl.collections.PatchCollection(edges1,cmap=plt.cm.Blues)
        pc.set_array(w)
        w=[(i-min(w))/(max(w)-min(w)) for i in w]
          #标准化为0-1之间
        for i in range(self.G.number_of_edges()):
            edges1[i].set_alpha(w[i])
        networkx.draw_networkx_labels(self.G,self.pos,font_size=font_size)
        networkx.draw_networkx_nodes(self.G,self.pos,node_size=node_size,node_color='#3697D6')
        if_new_units_connected=True
        try:
            networkx.draw_networkx_nodes(self.G,self.pos,nodelist=entry_node_list,node_size=node_size,node_color='r')
        except:
            if_new_units_connected=False
        ax=plt.gca()
        ax.set_axis_off()
        plt.colorbar(pc,ax=ax)
        plt.show()
        if not if_new_units_connected:
            print('New neuron units haven\'t been connected by others.')